import torch
from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial import KDTree
from sklearn.linear_model import LinearRegression
import umap
from ripser import ripser
from persim import plot_diagrams
from scipy.spatial.distance import pdist, squareform
from scipy.spatial.distance import cdist
from sklearn import svm
from sklearn import metrics
from sklearn.linear_model import LogisticRegression
import torch.nn as nn
import torch.optim as optim
from sklearn.neighbors import KernelDensity
from scipy.stats import entropy



def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

def main():

    A = [[0,1],[1,2],[2,2],[3,4],[3,5],[5,6],[6,7],[7,8],[8,9],[9,9]]
    B = [[0,1],[1,1],[2,2],[3,4],[4,5],[5,6],]
    C = [[0,1],[1,1],[2,2],[3,4],[4,5],[5,6],[6,7],[7,8]]
    # C = [[6,7],[7,8],[4,5],[5,6],[0,1],[1,1],[2,2],[3,4],]

    A = np.array(A)
    B = np.array(B)
    C = np.array(C)

    # # 将三个数组长度处理成一样的
    # max_length = max(A.shape[0], B.shape[0], C.shape[0])
    # if len(A) < max_length:
    #     last_element = A[-1]
    #     A = np.concatenate([A, np.repeat(last_element[np.newaxis,:], max_length - A.shape[0], axis=0)], axis=0)
    # if B.shape[0] < max_length:
    #     last_element = B[-1]
    #     B = np.concatenate([B, np.repeat(last_element[np.newaxis,:], max_length - B.shape[0], axis=0)], axis=0)
    # if len(C) < max_length:
    #     last_element = C[-1]
    #     C = np.concatenate([C, np.repeat(last_element[np.newaxis,:], max_length - C.shape[0], axis=0)], axis=0)

    def preprocess(trjs):
        max_length = max([trj.shape[0] for trj in trjs])
        for i in range(len(trjs)):
            trj = trjs[i]
            if trj.shape[0] < max_length:
                last_element = trj[-1]
                trjs[i] = np.concatenate([trj, np.repeat(last_element[np.newaxis,:], max_length - trj.shape[0], axis=0)], axis=0)
        return trjs
    
    @jax.jit
    def build_radiance_field(p):
        img = jnp.zeros((21, 21))
        top = 10
        bottom = 0
        effective_radius = 21*1.414
        px, py = p[0], p[1]

        x_coord_map = jnp.arange(img.shape[0])
        x_coord_map = jnp.repeat(x_coord_map[:, jnp.newaxis], img.shape[1], axis=1)

        y_coord_map = jnp.arange(img.shape[1])
        y_coord_map = jnp.repeat(y_coord_map[jnp.newaxis, :], img.shape[0], axis=0)
        
        dist = jnp.sqrt((x_coord_map - px)**2 + (y_coord_map - py)**2)
        img = jnp.where(dist < effective_radius,
                        (top - bottom) * (effective_radius - dist) / effective_radius + bottom,
                        img)
        return img
 
    @jax.jit
    def get_max_radiance_field(imgs):
        img = jnp.max(imgs, axis=0)
        return img
    
    @jax.jit
    def build_ridge(A):

        # 将 A 的第一个点对齐到 (10,10)
        A = A-A[0]+jnp.array([10,10])

        # 使用 jnp.map 来并行计算每个辐射场图像
        imgs = jax.vmap(build_radiance_field)(A)
        img = get_max_radiance_field(imgs)
        return img
    
    build_ridge_vmap = jax.vmap(build_ridge)


    trjs = preprocess([A, B, C])
    A, B, C = trjs[0], trjs[1], trjs[2]
    print("shape of A: ", A.shape)
    print("shape of B: ", B.shape)
    print("shape of C: ", C.shape)

    trjs = np.array(trjs)
    RIs = build_ridge_vmap(trjs)
    print("shape of RIs: ", RIs.shape)

    # 构建图形对象和子图
    fig, axs = plt.subplots(1, 2)

    # 绘制第一个子图
    img = build_ridge(A)
    img = np.flipud(img)
    axs[0].imshow(img)

    # 绘制第二个子图
    A = A - A[0] + np.array([10, 10])
    axs[1].plot(A[:, 0], A[:, 1], 'r')
    axs[1].set_xlim(0, 20)
    axs[1].set_ylim(0, 20)

    # 显示图形
    plt.show()

    # 将 img 作为三维曲面绘制出来
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    x = np.arange(0, 21, 1)
    y = np.arange(0, 21, 1)
    x, y = np.meshgrid(x, y)
    z = img
    ax.plot_surface(x, y, z, cmap='viridis')

    # 设置视角为俯视
    ax.view_init(elev=90, azim=0)  # 设置视角为从上方看下去
    plt.show()

    # img_A = build_ridge(A)
    # # 显示图像
    # plt.imshow(img_A, cmap='gray', interpolation='nearest')
    # plt.show()

    # img_B = build_ridge(B)
    # # 显示图像
    # plt.imshow(img_B, cmap='gray', interpolation='nearest')
    # plt.show()

    # t_start = time.time()
    # img_C = build_ridge(C)
    # t_end = time.time()
    # print("time cost: ", t_end - t_start)
    # # 显示图像
    # plt.imshow(img_C, cmap='gray', interpolation='nearest')
    # plt.show()

    # 计算 img_A/B/C 之间的距离
    dist_AB = np.linalg.norm(RIs[0] - RIs[1])
    dist_AC = np.linalg.norm(RIs[0] - RIs[2])
    dist_BC = np.linalg.norm(RIs[1] - RIs[2])

    print("dist_AB: ", dist_AB)
    print("dist_AC: ", dist_AC)
    print("dist_BC: ", dist_BC)




if __name__ == "__main__":
    main()
